{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example showing an analysis of a plate with a circular hole with a large deformation material model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using ForwardDiff\n",
    "using Tensors\n",
    "using JuAFEM\n",
    "using MAT\n",
    "using NLsolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constitutive law"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "immutable YeohMaterial{T}\n",
    "    λ::T\n",
    "    μ::T\n",
    "    c2::T\n",
    "    c3::T\n",
    "end\n",
    "\n",
    "function Ψ_Yeoh(C, mp::YeohMaterial)\n",
    "    μ, λ, c2, c3 = mp.μ, mp.λ, mp.c2, mp.c3\n",
    "    Ct = convert(Tensor{2, 2}, C)\n",
    "    J = sqrt(det(Ct))\n",
    "    Ic = trace(Ct) + 1\n",
    "    lnJ = log(J)\n",
    "    return μ/2 * (Ic - 3) + c2*(Ic - 3)^2 + c3 * (Ic - 3)^3 - μ*lnJ + λ/2 * lnJ^2\n",
    "end\n",
    "\n",
    "# Parameters used\n",
    "function get_material()\n",
    "    μ = 1.267e6\n",
    "    λ = 1.457e7\n",
    "    c2 = -μ/7.9\n",
    "    c3 = μ/41\n",
    "    return YeohMaterial(μ, λ, c2, c3)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Element internal forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Takes the initial + current configuration, the fe_value object and the material parameters and\n",
    "# computes the internal forces for the element.\n",
    "function intf_element{T,Q,dim}(x::Vector{T}, X::Vector{Q}, fe_values::FEValues{dim}, \n",
    "                               material_parameters::YeohMaterial)\n",
    "    # Closures for Forward Diff\n",
    "    Ψ_Yeohh(C) = Ψ_Yeoh(C, material_parameters)\n",
    "    ∂Ψ_Yeoh∂C = ForwardDiff.gradient(Ψ_Yeohh)\n",
    "    S_Yeoh(C) = 2 * ∂Ψ_Yeoh∂C(C)\n",
    "    \n",
    "    # Reinterpret x and X as vectors of first order tensors\n",
    "    n_basefuncs = n_basefunctions(get_functionspace(fe_values))\n",
    "    @assert length(x) == length(X) == dim * n_basefuncs\n",
    "    X_vec = reinterpret(Vec{dim, Q}, X, (n_basefuncs,))\n",
    "    x_vec = reinterpret(Vec{dim, T}, x, (n_basefuncs,))\n",
    "    \n",
    "    reinit!(fe_values, X_vec)\n",
    "    \n",
    "    fe = [zero(Tensor{1, dim, T}) for i in 1:n_basefuncs]\n",
    "\n",
    "    for q_point in 1:length(JuAFEM.points(get_quadrule(fe_values)))\n",
    "        F = function_vector_gradient(fe_values, q_point, x_vec)\n",
    "        C = F' ⋅ F\n",
    "        S = Tensor{2, 2}(S_Yeoh(vec(C)))\n",
    "        P = F ⋅ S\n",
    "        for i in 1:n_basefuncs\n",
    "            fe[i] += P ⋅ shape_gradient(fe_values, q_point, i) * detJdV(fe_values, q_point)\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    return reinterpret(T, fe, (dim * n_basefuncs,))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "immutable CALMesh2D\n",
    "    coord::Matrix{Float64}\n",
    "    dof::Matrix{Int}\n",
    "    edof::Matrix{Int}\n",
    "    ex::Matrix{Float64}\n",
    "    ey::Matrix{Float64}\n",
    "end\n",
    "\n",
    "# Reads the data from the .mat file and stores it in CALMesh2D as well as returning\n",
    "# the dofs that are free/prescribed and fixed,\n",
    "function read_data()\n",
    "    vars = matread(\"cass2_mesh_data.mat\");\n",
    "    dof_fixed = convert(Vector{Int}, vec(vars[\"dof_fixed\"]))\n",
    "    Coord = vars[\"Coord\"]'\n",
    "    Ex, Ey = vars[\"Ex\"], vars[\"Ey\"]\n",
    "    dof_prescr = convert(Vector{Int}, vec(vars[\"dof_prescr\"]))\n",
    "    Edof = convert(Matrix{Int}, vars[\"Edof\"])'[2:end,:]\n",
    "    Dof = convert(Matrix{Int}, vars[\"Dof\"]')\n",
    "    dof_free = convert(Vector{Int}, vec(vars[\"dof_free\"]));\n",
    "    return CALMesh2D(Coord, Dof, Edof, Ex, Ey), dof_prescr, dof_fixed, dof_free\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global internal forces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assembles the global internal forces as well as the reaction forces\n",
    "# for the prescribed dofs\n",
    "function internal_forces!(fvec, X, x, mesh, material_parameters, dof_free, dof_fixed, fe_values, f_react)\n",
    "    f_full = zeros(length(mesh.dof))\n",
    "\n",
    "    fill!(fvec, 0.0)\n",
    "    for i in 1:size(mesh.edof, 2)\n",
    "        dofs = mesh.edof[:, i]\n",
    "        fe = intf_element(x[dofs], X[dofs], fe_values, material_parameters)\n",
    "        f_full[dofs] += fe\n",
    "    end\n",
    "    f_react[dof_fixed] = f_full[dof_fixed]\n",
    "    copy!(fvec, f_full[dof_free])\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global stiffness matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function grad!(fvec, X, x, mesh, material_parameters, dof_free, fe_values)\n",
    "    n_basefuncs = n_basefunctions(get_functionspace(fe_values))\n",
    "\n",
    "    assembler = start_assemble()\n",
    "    XX = zeros(2 * n_basefuncs)\n",
    "    intf(x) = intf_element(x, XX, fe_values, material_parameters)\n",
    "    grad! = ForwardDiff.jacobian(intf, mutates = true)\n",
    "    Ke = zeros(2*n_basefuncs, 2*n_basefuncs)\n",
    "    for i in 1:size(mesh.edof, 2)\n",
    "        dofs = mesh.edof[:, i]\n",
    "        copy!(XX, X[dofs])\n",
    "        grad!(Ke, x[dofs])\n",
    "        assemble(dofs, assembler, Ke)\n",
    "    end\n",
    "    K = end_assemble(assembler)\n",
    "    return K[dof_free, dof_free]\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VTK output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function vtkoutput(pvd, i, mesh, X, x, topology, f_react)\n",
    "    u = x - X\n",
    "    nnodes = div(length(mesh.dof), 2)\n",
    "    nrelem = size(mesh.edof, 2)\n",
    "\n",
    "    disp = u\n",
    "    disp = reshape(disp, (2, nnodes))\n",
    "    disp = [disp; zeros(nnodes)']\n",
    "\n",
    "    f_react = reshape(f_react, (2, nnodes))\n",
    "    f_react = [f_react; zeros(nnodes)']\n",
    "\n",
    "\n",
    "    vtkfile = vtk_grid(topology, mesh.coord, \"box_$i\")\n",
    "    vtk_point_data(vtkfile, disp, \"displacement\")\n",
    "    vtk_point_data(vtkfile, f_react, \"reaction_forces\")\n",
    "    collection_add_timestep(pvd, vtkfile, float(i))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function go()\n",
    "    # Get the data\n",
    "    mesh, dof_prescr, dof_fixed, dof_free = read_data()\n",
    "    \n",
    "    # Get the material\n",
    "    material_parameters = get_material()\n",
    "\n",
    "    topology = topologyxtr(mesh.edof, mesh.coord, mesh.dof, 3)\n",
    "    pvd = paraview_collection(\"box\")\n",
    "    \n",
    "    function_space = Lagrange{2, RefTetrahedron, 1}()\n",
    "    quad_rule = QuadratureRule(Dim{2}, RefTetrahedron(), 1)\n",
    "    fe_values = FEValues(Float64, quad_rule, function_space)\n",
    "    \n",
    "    # Initialize\n",
    "    X = vec(mesh.coord)\n",
    "    x = copy(X)\n",
    "    prev_x = copy(X)\n",
    "    f_react = zeros(X)\n",
    "   \n",
    "    end_displacement = 30\n",
    "    nsteps = 20\n",
    "    for i in 1:nsteps\n",
    "        # Set current config to correct value.\n",
    "        prev_x[dof_prescr] = X[dof_prescr] + i / nsteps * end_displacement\n",
    "        \n",
    "        # Newton guess\n",
    "        dx0 = zeros(length(dof_free))\n",
    "\n",
    "        function f!(dx, fvec)\n",
    "            copy!(x, prev_x)\n",
    "            x[dof_free] += dx\n",
    "            internal_forces!(fvec, X, x, mesh, material_parameters, dof_free, dof_fixed, fe_values, f_react)\n",
    "        end\n",
    "\n",
    "        function g!(dx, g)\n",
    "            fvec = zeros(length(dof_free))\n",
    "            copy!(x, prev_x)\n",
    "            x[dof_free] += dx\n",
    "            K = grad!(fvec, X, x, mesh, material_parameters, dof_free, fe_values)\n",
    "            copy!(g, K)\n",
    "        end\n",
    "\n",
    "        println(\"Timestep $i out of $nsteps\")\n",
    "        df = DifferentiableSparseMultivariateFunction(f!, g!)\n",
    "        res = nlsolve(df, dx0; ftol = 1e-6, iterations = 20, method=:newton, show_trace=true)\n",
    "        if !converged(res)\n",
    "            error(\"Global equation did not converge\")\n",
    "        end\n",
    "        dx_conv = res.zero::Vector{Float64}\n",
    "        # Update converged solution\n",
    "        prev_x[dof_free] += dx_conv\n",
    "        vtkoutput(pvd, i, mesh, X, x, topology, f_react)\n",
    "    end\n",
    "    vtk_save(pvd)\n",
    "    return\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "go()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "julia-0.6",
   "display_name": "Julia 0.6.0",
   "language": "julia"
  },
  "language_info": {
   "mimetype": "application/julia",
   "file_extension": ".jl",
   "version": "0.6.0",
   "name": "julia"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
